import { create } from "zustand";
import { devtools } from "zustand/middleware";
import { immer } from "zustand/middleware/immer";
import { IResponseMessage, IUserMessage, Message } from "@/lib/types";

type ChatState = {
  currChatId: string | null;
  currMessageNodeId: string | null;
  messageMapping: {
    [key: string]: Message;
  };
  messages: Message[];
};

type ChatAction = {
  resetMessages: () => void;
  appendMessage: (
    message: IUserMessage | IResponseMessage,
    parent?: string | null,
  ) => void;
  response: (
    data: {
      chunk?: string;
      source?: { title: string; url: string }[];
      recommend?: string[];
      error?: any;
    },
    done: boolean,
  ) => void;
};

const useChatStore = create<ChatState & ChatAction>()(
  devtools(
    immer((set, get, store) => {
      return {
        currChatId: null,
        currMessageNodeId: null,
        messageMapping: {},
        messages: [],
        resetMessages() {
          set((state) => {
            state.messageMapping = {};
            state.messages = [];
            state.currMessageNodeId = null;
          });
        },
        appendMessage(message, parentId) {
          set((state) => {
            const newMessage: Message = {
              id: message.id,
              message: message,
              children: [],
              parent: null,
            };
            if (typeof parentId === "undefined") {
              // append message via input
              const parent = state.messages[state.messages.length - 1];
              if (parent) {
                newMessage.parent = parent.id;
                parent.children.push(message.id);
                if (parent.message.role === "assistant") {
                  parent.message.recommend = [];
                }
                state.messageMapping[parent.id] = parent;
              }
              state.messages.push(newMessage);
            } else if (typeof parentId === "string") {
              // append message via regenerate
              newMessage.parent = parentId;
              const parent = state.messageMapping[parentId];
              parent.children.push(message.id);
              if (parent.message.role === "assistant") {
                parent.message.recommend = [];
              }
              state.messageMapping[parentId] = parent;
              state.messageMapping[message.id] = newMessage;
              const newMessages = [newMessage];
              let pid: string | null = parentId;
              while (pid) {
                newMessages.unshift(state.messageMapping[pid]);
                pid = state.messageMapping[pid].parent;
              }
              state.messages = newMessages;
            } else {
              // regenerate the first message
              // currently this block will not work
              // it works before...
              state.messages = [newMessage];
            }
            state.messageMapping[message.id] = newMessage;
            state.currMessageNodeId = message.id;
          });
        },
        response(data, done) {
          set((state) => {
            if (!state.currMessageNodeId) {
              return;
            }
            const responseMessage = state.messages[state.messages.length - 1];
            if (!responseMessage || responseMessage.message.role === "user") {
              return;
            }
            const { chunk, error, recommend, source } = data;
            if (chunk) {
              responseMessage.message.content += chunk;
            }
            if (error) {
              responseMessage.message.error = error;
            }
            if (recommend) {
              responseMessage.message.recommend = recommend;
            }
            if (source) {
              responseMessage.message.source = source;
            }
            responseMessage.message.done = done;
            state.messageMapping[responseMessage.id] = responseMessage;
          });
        },
      };
    }),
  ),
);

export default useChatStore;
